package edu.zju.gis.dbfg.model.statistic;

import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.datasource.DataReaderRDD;
import edu.zju.gis.dbfg.model.util.CrsPick;
import edu.zju.gis.dbfg.model.util.GeoItem;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.broadcast.Broadcast;
import org.geotools.referencing.CRS;
import org.locationtech.jts.geom.Geometry;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

import static edu.zju.gis.dbfg.model.base.Constant.DEFAULT_CRS;

public class SelectByAttribute {

    private final static Logger logger = LoggerFactory.getLogger(SelectByAttribute.class);

    public static void main(String[] args) throws Exception{
        //Setup environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf();
//        conf.setMaster("local[4]");
//        conf.setAppName("select");
        JavaSparkContext jsc = new JavaSparkContext(conf);

        //Setup model arguments
        logger.info("Setup arguments");
        SelectByAttributeArgs mArgs = SelectByAttributeArgs.initArgs(args);
        if(mArgs == null){
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(SelectByAttribute.class, "main()", msg, args);
        }

        //Transform CRS
        CoordinateReferenceSystem crs = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs()));
        CoordinateReferenceSystem targetCRS = DEFAULT_CRS();

        MathTransform transform = null;
        if(!crs.equals(targetCRS)){
            transform = CRS.findMathTransform(crs, targetCRS);
        }

        //Read data from datasource and map it to geoItem
        String input = mArgs.getInput();
        String type = mArgs.getType();
        String filterIndex = mArgs.getFilterIndex();
        String filters = mArgs.getFilters();
        List<String> filterList = Arrays.asList(filters.split(","));

        Broadcast<String> bcFilterIndex = jsc.broadcast(filterIndex);
        Broadcast<List<String>> bcFilterList = jsc.broadcast(filterList);

        JavaRDD<GeoItem> gi = null;

        if(type.equals("ES")){
            gi = DataReaderRDD.readToRDD2(type, jsc, input, null, null, null, transform);
        }else{
            String idIndex = mArgs.getIdIndex();
            gi = DataReaderRDD.readToRDD(type, jsc, input, idIndex, null, transform);
        }

        //Filter elements by filterIndex,filters
        JavaRDD<GeoItem> temp = gi.filter(new Function<GeoItem, Boolean>() {
            @Override
            public Boolean call(GeoItem input) throws Exception {
                String filterValue = String.valueOf(input.getProp().get(bcFilterIndex.getValue()));
                if(bcFilterList.getValue().contains(filterValue)) {
                    return true;
                }
                return false;
            }
        });

        String output = mArgs.getOutput();
        writeToFile(temp, output, true);
    }

    public static void writeToFile(JavaRDD<GeoItem> rdd, String filename, Boolean saveAttribute) throws IOException {
        List<GeoItem> featureResult = rdd.collect();
        if(featureResult.size()==0)
            return;
        File featureFile = new File(filename);
        if (featureFile.exists()) featureFile.delete();
        FileWriter featureWriter = new FileWriter(featureFile);
//        StringBuilder fieldName = new StringBuilder();
//        fieldName.append("fid\t");
//        for(String key : featureResult.get(0).getProp().keySet()){
//            fieldName.append(key).append("\t");
//        }
//        fieldName.append("shape\n");
//        featureWriter.write(fieldName.toString());
        int fid = 0;
        for (GeoItem feature : featureResult) {
            GeoItem gi = feature;
            StringBuilder attribute = new StringBuilder();
            if(saveAttribute)
            {
                Map<String, Object> props = gi.getProp();
                for(String key: props.keySet()){
                    attribute.append(props.get(key));
                    attribute.append("\t");
                }
            }
            Geometry geometry = gi.getGeometry();
            String wkt = geometry.toText();
            fid++;
            featureWriter.write(String.format("%s\t%s%s\n", fid, attribute.toString(), wkt));
        }
        featureWriter.flush();
        featureWriter.close();
    }

}
